import cv2
import numpy as np
from typing import Optional
from xinfer_yolo.detect.preprocess_cpy import preprocess_cpy


def preprocess(src: np.ndarray, dst: np.ndarray, temp: Optional[np.ndarray] = None) -> np.ndarray:
    """ 目标检测任务图像预处理

    :param src: 输入图像, 通道格式为HWC, 色彩格式为BGR, 类型为uint8
    :param dst: 输出图像，通道格式为CHW，色彩格式为RGB, 类型为float32
    :param temp: 保存中间计算结果的图像, 为None则自动生成, 通道格式为HWC, 尺寸跟dst一致, 色彩格式为BGR, 类型为uint8
    :return: 返回仿射变换的逆变换矩阵
    """
    # todo完善错误处理
    if len(src.shape) != 3:
        raise RuntimeError
    if len(dst.shape) != 3:
        raise RuntimeError
    if temp is None:
        temp = np.zeros(shape=dst.shape, dtype=np.uint8)
    if len(temp.shape) != 3:
        raise RuntimeError

    # 获取图像尺寸信息
    src_height, src_width, _ = src.shape
    _, dst_width, dst_height = dst.shape
    # 计算缩放尺寸
    scale = min(dst_width // src_width, dst_height // src_height)
    ox = (dst_width - scale * src_width) / 2
    oy = (dst_height - scale * src_height) / 2
    # 构建仿射变换矩阵
    transformation_matrix = np.array([
        [scale, 0, ox],
        [0, scale, oy]
    ])
    # 构建逆变换矩阵
    invert_transformation_matrix = cv2.invertAffineTransform(transformation_matrix)
    # 进行仿射变换
    cv2.warpAffine(
        src, transformation_matrix, (dst_width, dst_height),
        dst=temp, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(114, 114, 114)
    )
    # 对dst进行BGR转RGB, HWC转CHW, 归一化
    preprocess_cpy(temp, dst, dst_height, dst_width)
    # 返回逆变换矩阵
    return invert_transformation_matrix
